# coding=utf-8
import os
import torch
import numpy as np
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms, utils , datasets
import matplotlib.pyplot as plt
from PIL import Image,ImageDraw
#from skimage import io, transform

#数据集可以在
#http://mmlab.ie.cuhk.edu.hk/projects/CelebA.html
#下载

print ("sdfsdf")
'''
def show_landmark(image, landmark):
    #tensor为有四个维度的张量，去掉第0维张量
    image = image.numpy().squeeze(0)
    #此时的张量是Channel，Height，Width的存储顺序
    #更改为Height，Width，Channel的顺序
    image = np.transpose(image,(1,2,0))
    #显示
    plt.imshow(image)
    #print(landmark)
    #画图象
    plt.scatter(landmark[0,: , 0], landmark[0,: , 1], s=10, marker='.', c='r')
    #暂停 等待完成更新
    plt.pause(0.001)

class ColorAugment(object):
    def __init__(self, color_range=0.2) :
        self.range = color_range
    def __call__(self, sample):
        image = sample['image']
        landmark = sample['landmark']
        #调整伽马
        image = transforms.functional.adjust_gamma(image, 1)
        sat_factor = 1 + (self.range - self.range * 2 * np.random.rand())
        #调整饱和度
        image = transforms.functional.adjust_saturation(image,sat_factor)
        return {'image': image, 'landmark': landmark}

class RandomCrop(object):
    def __init__(self, output_size):
        assert isinstance(output_size, (int, tuple))
        if isinstance(output_size, int):
            self.output_size = (output_size, output_size)
        else:
            assert len(output_size) == 2
            self.output_size = output_size

    def __call__(self, sample):
        image, landmark = sample['image'], sample['landmark']
        h = image.height
        w = image.width
        new_h, new_w = self.output_size
        #随机竖直方向的起始裁剪坐标
        top = np.random.randint(0, h - new_h)
        #随机水平方向的起始裁剪坐标
        left = np.random.randint(0, w - new_w)
        #裁剪
        image = image.crop((left,top, left + new_w, top + new_h))
        #坐标偏移
        landmark = landmark - [left, top]

        return {'image': image, 'landmark': landmark}

class ToTensor(object):
    def __call__(self, sample):
        image, landmark = sample['image'], sample['landmark']
        #转换为tensor
        image = transforms.ToTensor()(image)
        return {'image': image,
                'landmark': torch.from_numpy(landmark)}

class CelebADataset(Dataset):
    def __init__(self, anno, root_dir, transform=None):
        self.root_dir = root_dir
        self.transform = transform
        #打开标注文件
        anno_file = open(anno,'r');
        #读入文件为list
        self.anno = anno_file.readlines()
        #删除前两行
        self.anno.pop(0)
        self.anno.pop(0)
        self.image_list = []
        self.landmark_list = []
        self.transform = transform
        # 将路径存入image_list，将label存入landmark_list
        for item in self.anno:
            subitem = item.strip().split()
            self.image_list.append(subitem[0])
            landmark = list(map(int, subitem[1:]))
            self.landmark_list.append(landmark)

    def __len__(self):
        return len(self.anno)

    def __getitem__(self, idx):
        #组合root路径和图像路径
        img_name = os.path.join(self.root_dir,self.image_list[idx])
        #读取图像
        image = Image.open(img_name)
        #获得对应的label
        landmark = self.landmark_list[idx]
        #转换为numpy格式
        landmark = np.array(landmark)
        #转换为reshape为(5,2)的矩阵
        landmark = landmark.astype('float').reshape(-1, 2)
        sample = {'image': image, 'landmark': landmark}

        if self.transform:
            sample = self.transform(sample)

        return sample
        

def dataset_main(args):

    #创建多个transforms的list
    composed_transforms = transforms.Compose( 
    [RandomCrop(150), \
    ColorAugment(), \
    ToTensor()])
    # 创建CelebA Dataset
    dataset = CelebADataset('Anno\list_landmarks_align_celeba.txt', 'img_align_celeba',composed_transforms)
    # 创建CelebA Dataset的dataloader
    dataloader = DataLoader(dataset, batch_size=1, shuffle=True)
    epoches = 1
    for epoc in range(0 , epoches):
        for iteration, sample in enumerate(dataloader):
            #图像
            image = sample['image']
            #关键点
            landmark = sample['landmark']
            #print(landmark)
            #创建窗口
            plt.figure()
            #画图
            show_landmark(image,landmark)
            #显示
            plt.show()

if __name__ == '__main__':
    args = 0
    dataset_main(args)

'''